"""
离散化模块

本模块实现了双流体模型的数值离散化，包括：
1. 交错网格的实现
2. 各守恒方程的离散化
3. 离散矩阵的构建

理论基础：
- 交错网格法
- 有限体积法
- 离散化技术

作者：[您的名字]
日期：[创建日期]
"""

import numpy as np
from typing import Dict, Tuple, List
from dataclasses import dataclass
from ..physics.two_fluid_model import FluidState
from ..utils.logger import get_logger

@dataclass
class GridCell:
    """网格单元数据类"""
    # 几何参数
    dx: float  # x方向长度
    dy: float  # y方向长度
    dz: float  # z方向长度
    volume: float  # 体积
    
    # 面积
    Ax: float  # x方向面积
    Ay: float  # y方向面积
    Az: float  # z方向面积
    
    # 相邻单元索引
    neighbors: Dict[str, int]  # 相邻单元的索引

class StaggeredGrid:
    """
    交错网格类
    
    实现了交错网格的构建和管理，包括：
    - 标量网格（压力、密度等）
    - 矢量网格（速度）
    """
    
    def __init__(self, config: dict):
        """
        初始化交错网格
        
        参数:
            config: dict, 包含网格参数的配置字典
        """
        self.logger = get_logger()
        self.config = config
        
        # 网格参数
        self.nx = config['grid']['nx']  # x方向节点数
        self.ny = config['grid']['ny']  # y方向节点数
        self.nz = config['grid']['nz']  # z方向节点数
        
        self.dx = config['grid']['dx']  # x方向网格间距
        self.dy = config['grid']['dy']  # y方向网格间距
        self.dz = config['grid']['dz']  # z方向网格间距
        
        # 初始化网格
        self.scalar_cells = []  # 标量网格单元
        self.velocity_cells_x = []  # x方向速度网格单元
        self.velocity_cells_y = []  # y方向速度网格单元
        self.velocity_cells_z = []  # z方向速度网格单元
        
        self._initialize_grid()
        
    def _initialize_grid(self):
        """初始化交错网格"""
        # 初始化标量网格
        for k in range(self.nz):
            for j in range(self.ny):
                for i in range(self.nx):
                    # 计算体积和面积
                    volume = self.dx * self.dy * self.dz
                    Ax = self.dy * self.dz
                    Ay = self.dx * self.dz
                    Az = self.dx * self.dy
                    
                    # 计算相邻单元索引
                    neighbors = {
                        'east': i + 1 if i < self.nx - 1 else -1,
                        'west': i - 1 if i > 0 else -1,
                        'north': j + 1 if j < self.ny - 1 else -1,
                        'south': j - 1 if j > 0 else -1,
                        'top': k + 1 if k < self.nz - 1 else -1,
                        'bottom': k - 1 if k > 0 else -1
                    }
                    
                    cell = GridCell(
                        dx=self.dx,
                        dy=self.dy,
                        dz=self.dz,
                        volume=volume,
                        Ax=Ax,
                        Ay=Ay,
                        Az=Az,
                        neighbors=neighbors
                    )
                    
                    self.scalar_cells.append(cell)
        
        # 初始化速度网格（交错布置）
        # x方向速度网格
        for k in range(self.nz):
            for j in range(self.ny):
                for i in range(self.nx + 1):  # 多一个节点
                    volume = self.dx * self.dy * self.dz
                    Ax = self.dy * self.dz
                    cell = GridCell(
                        dx=self.dx,
                        dy=self.dy,
                        dz=self.dz,
                        volume=volume,
                        Ax=Ax,
                        Ay=0.0,
                        Az=0.0,
                        neighbors={}
                    )
                    self.velocity_cells_x.append(cell)
        
        # y方向速度网格
        for k in range(self.nz):
            for j in range(self.ny + 1):  # 多一个节点
                for i in range(self.nx):
                    volume = self.dx * self.dy * self.dz
                    Ay = self.dx * self.dz
                    cell = GridCell(
                        dx=self.dx,
                        dy=self.dy,
                        dz=self.dz,
                        volume=volume,
                        Ax=0.0,
                        Ay=Ay,
                        Az=0.0,
                        neighbors={}
                    )
                    self.velocity_cells_y.append(cell)
        
        # z方向速度网格
        for k in range(self.nz + 1):  # 多一个节点
            for j in range(self.ny):
                for i in range(self.nx):
                    volume = self.dx * self.dy * self.dz
                    Az = self.dx * self.dy
                    cell = GridCell(
                        dx=self.dx,
                        dy=self.dy,
                        dz=self.dz,
                        volume=volume,
                        Ax=0.0,
                        Ay=0.0,
                        Az=Az,
                        neighbors={}
                    )
                    self.velocity_cells_z.append(cell)

class Discretization:
    """
    离散化类
    
    实现了双流体模型各守恒方程的离散化，包括：
    - 质量守恒方程
    - 动量守恒方程（轴向和横向）
    - 能量守恒方程
    """
    
    def __init__(self, grid: StaggeredGrid, config: dict):
        """
        初始化离散化器
        
        参数:
            grid: StaggeredGrid, 交错网格对象
            config: dict, 配置字典
        """
        self.logger = get_logger()
        self.grid = grid
        self.config = config
        
        # 时间步长
        self.dt = config['numerical']['time_step']
        
    def discretize_mass_conservation(
        self,
        k_phase: FluidState,
        prev_state: FluidState,
        gamma: float,
        T_mk: float,
        cell_idx: int
    ) -> float:
        """
        离散化质量守恒方程
        
        参数:
            k_phase: 当前时刻k相状态
            prev_state: 上一时刻状态
            gamma: 相变质量流率
            T_mk: 湍流混合质量流率
            cell_idx: 网格单元索引
            
        返回:
            float: 质量守恒方程的残差
        """
        cell = self.grid.scalar_cells[cell_idx]
        
        # 时间导数项
        d_rho_dt = (k_phase.alpha * k_phase.rho - prev_state.alpha * prev_state.rho) * \
                   cell.volume / self.dt
        
        # 对流项（三个方向）
        flux_x = self._calculate_mass_flux_x(k_phase, cell_idx)
        flux_y = self._calculate_mass_flux_y(k_phase, cell_idx)
        flux_z = self._calculate_mass_flux_z(k_phase, cell_idx)
        
        # 源项
        source = (gamma + T_mk) * cell.volume
        
        return d_rho_dt + flux_x + flux_y + flux_z - source
    
    def discretize_axial_momentum(
        self,
        k_phase: FluidState,
        prev_state: FluidState,
        P: float,
        F_wk: float,
        F_ik: float,
        M_tk: float,
        cell_idx: int
    ) -> float:
        """
        离散化轴向动量守恒方程
        
        参数:
            k_phase: 当前时刻k相状态
            prev_state: 上一时刻状态
            P: 压力
            F_wk: 壁面剪切力
            F_ik: 界���剪切力
            M_tk: 湍流混合动量传递
            cell_idx: 网格单元索引
            
        返回:
            float: 轴向动量守恒方程的残差
        """
        cell = self.grid.velocity_cells_x[cell_idx]
        
        # 时间导数项
        d_mom_dt = (k_phase.alpha * k_phase.rho * k_phase.u - 
                    prev_state.alpha * prev_state.rho * prev_state.u) * \
                   cell.volume / self.dt
        
        # 动量通量项
        flux_x = self._calculate_momentum_flux_x(k_phase, cell_idx)
        flux_y = self._calculate_momentum_flux_y(k_phase, cell_idx)
        flux_z = self._calculate_momentum_flux_z(k_phase, cell_idx)
        
        # 压力梯度项
        dp_dx = self._calculate_pressure_gradient_x(P, cell_idx)
        pressure_term = -k_phase.alpha * dp_dx * cell.volume
        
        # 重力项
        gravity_term = -k_phase.alpha * k_phase.rho * self.config['physics']['gravity'] * \
                      cell.volume
        
        # 源项
        source = (-F_wk + F_ik + M_tk) * cell.volume
        
        return d_mom_dt + flux_x + flux_y + flux_z + pressure_term + gravity_term - source
    
    def discretize_lateral_momentum(
        self,
        k_phase: FluidState,
        prev_state: FluidState,
        P: float,
        F_wk: float,
        F_ik: float,
        M_tk: float,
        cell_idx: int
    ) -> float:
        """
        离散化横向动量守恒方程
        
        参数:
            k_phase: 当前时刻k相状态
            prev_state: 上一时刻状态
            P: 压力
            F_wk: 壁面剪切力
            F_ik: 界面剪切力
            M_tk: 湍流混合动量传递
            cell_idx: 网格单元索引
            
        返回:
            float: 横向动量守恒方程的残差
        """
        cell = self.grid.velocity_cells_y[cell_idx]
        
        # 时间导数项
        d_mom_dt = (k_phase.alpha * k_phase.rho * k_phase.v - 
                    prev_state.alpha * prev_state.rho * prev_state.v) * \
                   cell.volume / self.dt
        
        # 动量通量项（简化处理，忽略横向动量的正交方向通量）
        flux_x = self._calculate_momentum_flux_x(k_phase, cell_idx)
        
        # 压力梯度项
        dp_dy = self._calculate_pressure_gradient_y(P, cell_idx)
        pressure_term = -k_phase.alpha * dp_dy * cell.volume
        
        # 源项
        source = (-F_wk + F_ik + M_tk) * cell.volume
        
        return d_mom_dt + flux_x + pressure_term - source
    
    def discretize_energy(
        self,
        k_phase: FluidState,
        prev_state: FluidState,
        P: float,
        prev_P: float,
        q_wk: float,
        gamma_h: float,
        T_kr: float,
        cell_idx: int
    ) -> float:
        """
        离散化能量守恒方程
        
        参数:
            k_phase: 当前时刻k相状态
            prev_state: 上一时刻状态
            P: 当前压力
            prev_P: 上一时刻压力
            q_wk: 壁面传热功率
            gamma_h: 相变焓传递
            T_kr: 湍流混合能量传递
            cell_idx: 网格单元索引
            
        返回:
            float: 能量守恒方程的残差
        """
        cell = self.grid.scalar_cells[cell_idx]
        
        # 时间导数项
        d_energy_dt = (k_phase.alpha * k_phase.rho * k_phase.h - 
                      prev_state.alpha * prev_state.rho * prev_state.h) * \
                     cell.volume / self.dt
        
        # 能量通量项
        flux_x = self._calculate_energy_flux_x(k_phase, cell_idx)
        flux_y = self._calculate_energy_flux_y(k_phase, cell_idx)
        flux_z = self._calculate_energy_flux_z(k_phase, cell_idx)
        
        # 压力功项
        dp_dt = (P - prev_P) / self.dt
        pressure_work = k_phase.alpha * dp_dt * cell.volume
        
        # 源项
        source = (q_wk + gamma_h + T_kr) * cell.volume
        
        return d_energy_dt + flux_x + flux_y + flux_z - pressure_work - source
    
    def _calculate_mass_flux_x(self, k_phase: FluidState, cell_idx: int) -> float:
        """计算x方向质量通量，使用一阶迎风格式"""
        cell = self.grid.scalar_cells[cell_idx]
        
        # 获取相邻单元索引
        i_east = cell.neighbors['east']
        i_west = cell.neighbors['west']
        
        # 计算界面通量
        flux = 0.0
        
        # 东界面
        if i_east >= 0:
            u_e = k_phase.u  # 界面速度
            if u_e > 0:  # 正向流动
                flux += cell.Ax * k_phase.alpha * k_phase.rho * u_e
            else:  # 反向流动
                flux += cell.Ax * k_phase.alpha * k_phase.rho * u_e
        
        # 西界面
        if i_west >= 0:
            u_w = k_phase.u  # 界面速度
            if u_w > 0:  # 正向流动
                flux -= cell.Ax * k_phase.alpha * k_phase.rho * u_w
            else:  # 反向流动
                flux -= cell.Ax * k_phase.alpha * k_phase.rho * u_w
        
        return flux

    def _calculate_mass_flux_y(self, k_phase: FluidState, cell_idx: int) -> float:
        """计算y方向质量通量，使用一阶迎风格式"""
        cell = self.grid.scalar_cells[cell_idx]
        
        # 获取相邻单元索引
        i_north = cell.neighbors['north']
        i_south = cell.neighbors['south']
        
        # 计算界面通量
        flux = 0.0
        
        # 北界面
        if i_north >= 0:
            v_n = k_phase.v  # 界面速度
            if v_n > 0:  # 正向流动
                flux += cell.Ay * k_phase.alpha * k_phase.rho * v_n
            else:  # 反向流动
                flux += cell.Ay * k_phase.alpha * k_phase.rho * v_n
        
        # 南界面
        if i_south >= 0:
            v_s = k_phase.v  # 界面速度
            if v_s > 0:  # 正向流动
                flux -= cell.Ay * k_phase.alpha * k_phase.rho * v_s
            else:  # 反向流动
                flux -= cell.Ay * k_phase.alpha * k_phase.rho * v_s
        
        return flux

    def _calculate_mass_flux_z(self, k_phase: FluidState, cell_idx: int) -> float:
        """计算z方向质量通量，使用一阶迎风格式"""
        cell = self.grid.scalar_cells[cell_idx]
        
        # 获取相邻单元索引
        i_top = cell.neighbors['top']
        i_bottom = cell.neighbors['bottom']
        
        # 计算界面通量
        flux = 0.0
        
        # 顶部界面
        if i_top >= 0:
            w_t = k_phase.w  # 界面速度
            if w_t > 0:  # 正向流动
                flux += cell.Az * k_phase.alpha * k_phase.rho * w_t
            else:  # 反向流动
                flux += cell.Az * k_phase.alpha * k_phase.rho * w_t
        
        # 底部界面
        if i_bottom >= 0:
            w_b = k_phase.w  # 界面速度
            if w_b > 0:  # 正向流动
                flux -= cell.Az * k_phase.alpha * k_phase.rho * w_b
            else:  # 反向流动
                flux -= cell.Az * k_phase.alpha * k_phase.rho * w_b
        
        return flux

    def _calculate_momentum_flux_x(self, k_phase: FluidState, cell_idx: int) -> float:
        """计算x方向动量通量，使用一阶迎风格式"""
        cell = self.grid.velocity_cells_x[cell_idx]
        
        # 获取相邻单元索引
        i_east = cell_idx + 1 if cell_idx < len(self.grid.velocity_cells_x) - 1 else -1
        i_west = cell_idx - 1 if cell_idx > 0 else -1
        
        # 计算界面通量
        flux = 0.0
        
        # 东界面
        if i_east >= 0:
            u_e = k_phase.u  # 界面速度
            if u_e > 0:  # 正向流动
                flux += cell.Ax * k_phase.alpha * k_phase.rho * u_e * u_e
            else:  # 反向流动
                flux += cell.Ax * k_phase.alpha * k_phase.rho * u_e * u_e
        
        # 西界面
        if i_west >= 0:
            u_w = k_phase.u  # 界面速度
            if u_w > 0:  # 正向流动
                flux -= cell.Ax * k_phase.alpha * k_phase.rho * u_w * u_w
            else:  # 反向流动
                flux -= cell.Ax * k_phase.alpha * k_phase.rho * u_w * u_w
        
        return flux

    def _calculate_momentum_flux_y(self, k_phase: FluidState, cell_idx: int) -> float:
        """计算y方向动量通量，使用一阶迎风格式"""
        cell = self.grid.velocity_cells_y[cell_idx]
        
        # 获取相邻单元索引
        i_north = cell_idx + 1 if cell_idx < len(self.grid.velocity_cells_y) - 1 else -1
        i_south = cell_idx - 1 if cell_idx > 0 else -1
        
        # 计算界面通量
        flux = 0.0
        
        # 北界面
        if i_north >= 0:
            v_n = k_phase.v  # 界面速度
            if v_n > 0:  # 正向流动
                flux += cell.Ay * k_phase.alpha * k_phase.rho * v_n * v_n
            else:  # 反向流动
                flux += cell.Ay * k_phase.alpha * k_phase.rho * v_n * v_n
        
        # 南界面
        if i_south >= 0:
            v_s = k_phase.v  # 界面速度
            if v_s > 0:  # 正向流动
                flux -= cell.Ay * k_phase.alpha * k_phase.rho * v_s * v_s
            else:  # 反向流动
                flux -= cell.Ay * k_phase.alpha * k_phase.rho * v_s * v_s
        
        return flux

    def _calculate_momentum_flux_z(self, k_phase: FluidState, cell_idx: int) -> float:
        """计算z方向动量通量，使用一阶迎风格式"""
        cell = self.grid.velocity_cells_z[cell_idx]
        
        # 获取相邻单元索引
        i_top = cell_idx + 1 if cell_idx < len(self.grid.velocity_cells_z) - 1 else -1
        i_bottom = cell_idx - 1 if cell_idx > 0 else -1
        
        # 计算界面通量
        flux = 0.0
        
        # 顶部界面
        if i_top >= 0:
            w_t = k_phase.w  # 界面速度
            if w_t > 0:  # 正向流动
                flux += cell.Az * k_phase.alpha * k_phase.rho * w_t * w_t
            else:  # 反向流动
                flux += cell.Az * k_phase.alpha * k_phase.rho * w_t * w_t
        
        # 底部界面
        if i_bottom >= 0:
            w_b = k_phase.w  # 界面速度
            if w_b > 0:  # 正向流动
                flux -= cell.Az * k_phase.alpha * k_phase.rho * w_b * w_b
            else:  # 反向流动
                flux -= cell.Az * k_phase.alpha * k_phase.rho * w_b * w_b
        
        return flux

    def _calculate_energy_flux_x(self, k_phase: FluidState, cell_idx: int) -> float:
        """计算x方向能量通量，使用一阶迎风格式"""
        cell = self.grid.scalar_cells[cell_idx]
        
        # 获取相邻单元索引
        i_east = cell.neighbors['east']
        i_west = cell.neighbors['west']
        
        # 计算界面通量
        flux = 0.0
        
        # 东界面
        if i_east >= 0:
            u_e = k_phase.u  # 界面速度
            if u_e > 0:  # 正向流动
                flux += cell.Ax * k_phase.alpha * k_phase.rho * k_phase.h * u_e
            else:  # 反向流动
                flux += cell.Ax * k_phase.alpha * k_phase.rho * k_phase.h * u_e
        
        # 西界面
        if i_west >= 0:
            u_w = k_phase.u  # 界面速度
            if u_w > 0:  # 正向流动
                flux -= cell.Ax * k_phase.alpha * k_phase.rho * k_phase.h * u_w
            else:  # 反向流动
                flux -= cell.Ax * k_phase.alpha * k_phase.rho * k_phase.h * u_w
        
        return flux

    def _calculate_energy_flux_y(self, k_phase: FluidState, cell_idx: int) -> float:
        """计算y方向能量通量，使用一阶迎风格式"""
        cell = self.grid.scalar_cells[cell_idx]
        
        # 获取相邻单元索引
        i_north = cell.neighbors['north']
        i_south = cell.neighbors['south']
        
        # 计算界面通量
        flux = 0.0
        
        # 北界面
        if i_north >= 0:
            v_n = k_phase.v  # 界面速度
            if v_n > 0:  # 正向流动
                flux += cell.Ay * k_phase.alpha * k_phase.rho * k_phase.h * v_n
            else:  # 反向流动
                flux += cell.Ay * k_phase.alpha * k_phase.rho * k_phase.h * v_n
        
        # 南界面
        if i_south >= 0:
            v_s = k_phase.v  # 界面速度
            if v_s > 0:  # 正向流动
                flux -= cell.Ay * k_phase.alpha * k_phase.rho * k_phase.h * v_s
            else:  # 反向流动
                flux -= cell.Ay * k_phase.alpha * k_phase.rho * k_phase.h * v_s
        
        return flux

    def _calculate_energy_flux_z(self, k_phase: FluidState, cell_idx: int) -> float:
        """计算z方向能量通量，使用一阶迎风格式"""
        cell = self.grid.scalar_cells[cell_idx]
        
        # 获取相邻单元索引
        i_top = cell.neighbors['top']
        i_bottom = cell.neighbors['bottom']
        
        # 计算界面通量
        flux = 0.0
        
        # 顶部界面
        if i_top >= 0:
            w_t = k_phase.w  # 界面速度
            if w_t > 0:  # 正向流动
                flux += cell.Az * k_phase.alpha * k_phase.rho * k_phase.h * w_t
            else:  # 反向流动
                flux += cell.Az * k_phase.alpha * k_phase.rho * k_phase.h * w_t
        
        # 底部界面
        if i_bottom >= 0:
            w_b = k_phase.w  # 界面速度
            if w_b > 0:  # 正向流动
                flux -= cell.Az * k_phase.alpha * k_phase.rho * k_phase.h * w_b
            else:  # 反向流动
                flux -= cell.Az * k_phase.alpha * k_phase.rho * k_phase.h * w_b
        
        return flux

    def _calculate_pressure_gradient_x(self, P: float, cell_idx: int) -> float:
        """计算x方向压力梯度，使用中心差分"""
        cell = self.grid.scalar_cells[cell_idx]
        
        # 获取相邻单元索引
        i_east = cell.neighbors['east']
        i_west = cell.neighbors['west']
        
        if i_east >= 0 and i_west >= 0:
            # 中心差分
            return (P[i_east] - P[i_west]) / (2 * cell.dx)
        elif i_east >= 0:
            # 前向差分
            return (P[i_east] - P) / cell.dx
        elif i_west >= 0:
            # 后向差分
            return (P - P[i_west]) / cell.dx
        else:
            return 0.0

    def _calculate_pressure_gradient_y(self, P: float, cell_idx: int) -> float:
        """计算y方向压力梯度，使用中心差分"""
        cell = self.grid.scalar_cells[cell_idx]
        
        # 获取相邻单元索引
        i_north = cell.neighbors['north']
        i_south = cell.neighbors['south']
        
        if i_north >= 0 and i_south >= 0:
            # 中心差分
            return (P[i_north] - P[i_south]) / (2 * cell.dy)
        elif i_north >= 0:
            # 前向差分
            return (P[i_north] - P) / cell.dy
        elif i_south >= 0:
            # 后向差分
            return (P - P[i_south]) / cell.dy
        else:
            return 0.0 